//
// Created by subject on 2025/2/21.
//

#include <JavaManager.h>

std::mutex g_destroyMutex;

// 全局变量，用于存储初始化后的类和方法引用
static jclass g_lastHiddenLayerClass = nullptr;
static jclass g_resultClass = nullptr;

static jclass g_llmCallStateClass = nullptr;
static jclass g_objectClass = nullptr;

static jmethodID g_lastHiddenLayerConstructor = nullptr;
static jmethodID g_resultConstructor = nullptr;
static jmethodID g_objectConstructor = nullptr;

static jmethodID g_onGetResultMethod = nullptr;
static jmethodID g_dsFinish = nullptr;

static jobject g_dsManagerInstance = nullptr;

bool JavaManager::initJavaMethod(JNIEnv *env) {
    // 初始化 RKLLMResultLastHiddenLayer 类和构造方法

//    a1. 获取 DsManager 类的引用
    jclass dsManagerClass = findClass(env,"com/taiic/deepseek/DsManager");
//    a2. 获取 DsManager 的 GetInstance 方法ID
    jmethodID getInstanceMethod = getStaticMethodID(env, dsManagerClass, "GetInstance", "()Lcom/taiic/deepseek/DsManager;");
//    a3. 调用 getInstance 方法， 获取 DsManager 的单例实例
    jobject dsManagerInstance = env->CallStaticObjectMethod(dsManagerClass, getInstanceMethod);

//    b1. 获取 RKLLMResultLastHiddenLayer 类的引用
    jclass lastHiddenLayerClass = findClass(env, "com/taiic/deepseek/rkllm/RKLLMResultLastHiddenLayer");
//    b2. 获取 RKLLMResultLastHiddenLayer 的 init 方法ID
    g_lastHiddenLayerConstructor = getMethodID(env, lastHiddenLayerClass, "<init>", "([FII)V");

//    c1. 获取 RKLLMResult 类的引用
    jclass resultClass = findClass(env, "com/taiic/deepseek/rkllm/RKLLMResult");
//    c2. 获取 RKLLMResult 的 init 方法ID
    g_resultConstructor = getMethodID(env, resultClass, "<init>", "(Ljava/lang/String;ILcom/taiic/deepseek/rkllm/RKLLMResultLastHiddenLayer;)V");

//    d1. 获取 LLMCallState 类的引用
    jclass llmCallStateClass = findClass(env,"com/taiic/deepseek/rkllm/LLMCallState");

//    e1. 找到 java.lang.Object 类
    jclass objectClass = findClass(env,"java/lang/Object");
//    e2. 获取 Object 的构造方法
    g_objectConstructor = getMethodID(env, objectClass, "<init>", "()V");

//    a4. 将局部引用转为全局引用
    if (dsManagerInstance == nullptr) {
        LOGE("Failed to get DsManager instance.");
        safeDeleteLocalRef(env, dsManagerClass);
        safeDeleteLocalRef(env, dsManagerInstance);
        goto error;
    }
    g_dsManagerInstance = env->NewGlobalRef(dsManagerInstance);
    if(g_dsManagerInstance == nullptr){
        LOGE("Failed to create global reference for DsManager instance\n");
        safeDeleteLocalRef(env, dsManagerClass);
        safeDeleteLocalRef(env, dsManagerInstance);
        goto error;
    }
//    a5.1 获取 onGetResult 的方法ID
    g_onGetResultMethod = getMethodID(env, dsManagerClass, "onGetResult", "(Lcom/taiic/deepseek/rkllm/RKLLMResult;Ljava/lang/Object;Lcom/taiic/deepseek/rkllm/LLMCallState;)V");
    if (g_onGetResultMethod == nullptr) goto error;
//    a5.2 获取 onFinsh 的方法ID
    g_dsFinish = getMethodID(env, dsManagerClass, "onFinsh", "()V");
    if(g_dsFinish == nullptr) goto error;
//    a6. 删除局部引用
    if (dsManagerClass == nullptr) goto error;
    safeDeleteLocalRef(env, dsManagerClass);
    safeDeleteLocalRef(env, dsManagerInstance);

//    b3. 将局部引用转为全局引用
    g_lastHiddenLayerClass = (jclass)env->NewGlobalRef(lastHiddenLayerClass);
    if(g_lastHiddenLayerClass == nullptr){
        LOGE("Failed to create global reference for RKLLMResultLastHiddenLayer class.\n");
        safeDeleteLocalRef(env, lastHiddenLayerClass);
        goto error;
    }
//    b4. 检查全局引用
    if (g_lastHiddenLayerClass == nullptr) goto error;
    if (g_lastHiddenLayerConstructor == nullptr) goto error;
    
//    c3. 将局部引用转为全局引用
    g_resultClass = (jclass)env->NewGlobalRef(resultClass);
    if(g_resultClass == nullptr){
        LOGE("Failed to create global reference for RKLLMResult class.\n");
        safeDeleteLocalRef(env, resultClass);
        goto error;
    }
//    c4. 检查全局引用
    if (g_resultConstructor == nullptr) goto error;

//    d2. 将局部引用转为全局引用
    g_llmCallStateClass = (jclass)env->NewGlobalRef(llmCallStateClass);
    if(g_llmCallStateClass == nullptr){
        LOGE("Failed to create global reference for LLMCallState class.\n");
        safeDeleteLocalRef(env, llmCallStateClass);
        goto error;
    }

//    e3. 将局部引用转为全局引用
    g_objectClass = (jclass)env->NewGlobalRef(objectClass);
    if (g_objectClass == nullptr)
    {
        LOGE("Failed to create global reference for Object class.\n");
        safeDeleteLocalRef(env, objectClass);
        goto error;
    }
//    e4. 检查全局引用
    if (g_objectConstructor == nullptr) goto error;

    return true;

error:
    destoryJavaMethod(env);
    return false;

}

void JavaManager::destoryJavaMethod(JNIEnv *env) {
    std::lock_guard<std::mutex> lock(g_destroyMutex);

    safeDeleteGlobalRef(env, g_lastHiddenLayerClass);
    safeDeleteGlobalRef(env, g_resultClass);
    safeDeleteGlobalRef(env, g_llmCallStateClass);
    safeDeleteGlobalRef(env, g_objectClass);

    safeDeleteGlobalRef(env, g_dsManagerInstance);
}

bool
JavaManager::callLlmResult(JNIEnv *env, RKLLMResult *result, void *userdata, LLMCallState state) {
    if (g_dsManagerInstance == nullptr || g_onGetResultMethod == nullptr) {
        LOGE("DsManager instance or onGetResult method is not initialized.");
        return false;
    }

//    Convert RKLLMResultLastHiddenLayer to Java object
//    Create float[] for hidden_states
    int hiddenStatesSize = result->last_hidden_layer.num_tokens * result->last_hidden_layer.embd_size;
    jfloatArray hiddenStatesArray = env->NewFloatArray(hiddenStatesSize);
    if (hiddenStatesArray == nullptr) {
        LOGE("Failed to allocate float array.");
        return false;
    }
    env->SetFloatArrayRegion(hiddenStatesArray, 0, hiddenStatesSize, result->last_hidden_layer.hidden_states);
//    Create RKLLMResultLastHiddenLayer object
    jobject lastHiddenLayerObject = env->NewObject(
            g_lastHiddenLayerClass,
            g_lastHiddenLayerConstructor,
            hiddenStatesArray,
            result->last_hidden_layer.embd_size,
            result->last_hidden_layer.num_tokens
    );
    safeDeleteLocalRef(env, hiddenStatesArray);

    if (lastHiddenLayerObject == nullptr) {
        LOGE("Failed to create RKLLMResultLastHiddenLayer object.");
        return false;
    }
//    Convert RKLLMResult to Java object
//    Create String for text
    jstring text = env->NewStringUTF(result->text);
    if (text == nullptr) {
        LOGE("Failed to create string for RKLLMResult.");
        safeDeleteLocalRef(env, lastHiddenLayerObject); // Clean up lastHiddenLayerObject
        return false;
    }
//    Create RKLLMResult object
    jobject rkllmResultObject = env->NewObject(g_resultClass, g_resultConstructor, text, result->token_id, lastHiddenLayerObject);
    safeDeleteLocalRef(env, text);
    safeDeleteLocalRef(env, lastHiddenLayerObject);
    if (rkllmResultObject == nullptr) {
        LOGE("Failed to create RKLLMResult object.");
        return false;
    }
//    Convert LLMCallState to Java enum
    jmethodID valueOfMethod = env->GetStaticMethodID(g_llmCallStateClass, "valueOf", "(Ljava/lang/String;)Lcom/taiic/deepseek/rkllm/LLMCallState;");
    if (valueOfMethod == nullptr) {
        LOGE("Failed to find valueOf method for LLMCallState.");
        safeDeleteLocalRef(env, rkllmResultObject);
        return false;
    }
//    Map C++ enum to Java enum name
    const char* stateName = nullptr;
    switch (state) {
        case RKLLM_RUN_NORMAL: stateName = "RKLLM_RUN_NORMAL"; break;
        case RKLLM_RUN_WAITING: stateName = "RKLLM_RUN_WAITING"; break;
        case RKLLM_RUN_FINISH: stateName = "RKLLM_RUN_FINISH"; break;
        case RKLLM_RUN_ERROR: stateName = "RKLLM_RUN_ERROR"; break;
        case RKLLM_RUN_GET_LAST_HIDDEN_LAYER: stateName = "RKLLM_RUN_GET_LAST_HIDDEN_LAYER"; break;
        default: stateName = "RKLLM_RUN_ERROR"; break;
    }
    jstring stateString = env->NewStringUTF(stateName);
    jobject llmCallStateObject = env->CallStaticObjectMethod(g_llmCallStateClass, valueOfMethod, stateString);
    safeDeleteLocalRef(env, stateString);
    if (llmCallStateObject == nullptr) {
        LOGE("Failed to create LLMCallState object.");
        safeDeleteLocalRef(env, rkllmResultObject); // Clean up rkllmResultObject
        return false;
    }

//  创建一个空的 Object 实例
    static jobject emptyObject = nullptr;
    if (emptyObject == nullptr) {
        emptyObject = env->NewGlobalRef(env->NewObject(g_objectClass, g_objectConstructor));
        if (emptyObject == nullptr) {
            LOGE("Failed to create an empty Object.");
            safeDeleteLocalRef(env, rkllmResultObject); // Clean up rkllmResultObject
            safeDeleteLocalRef(env, llmCallStateObject); // Clean up llmCallStateObject
            return false;
        }
    }

//    Call onGetResult method
    env->CallVoidMethod(g_dsManagerInstance, g_onGetResultMethod, rkllmResultObject, emptyObject, llmCallStateObject);

    // Clean up local references
    safeDeleteLocalRef(env, rkllmResultObject);
    safeDeleteLocalRef(env, llmCallStateObject);
    return true;
}

// private method
jclass JavaManager::findClass(JNIEnv *env, const char *className) {
    jclass clazz = env->FindClass(className);
    if (clazz == nullptr) {
        LOGE("Failed to find class: %s", className);
    }
    return clazz;
}

jmethodID
JavaManager::getMethodID(JNIEnv *env, jclass clazz, const char *methodName, const char *methodSig) {
    jmethodID methodID = env->GetMethodID(clazz, methodName, methodSig);
    if (methodID == nullptr) {
        LOGE("Failed to find method: %s with signature: %s", methodName, methodSig);
    }
    return methodID;
}

jmethodID JavaManager::getStaticMethodID(JNIEnv *env, jclass clazz, const char *methodName,
                                         const char *methodSig) {
    jmethodID methodID = env->GetStaticMethodID(clazz, methodName, methodSig);
    if (methodID == nullptr) {
        LOGE("Failed to find static method: %s with signature: %s", methodName, methodSig);
    }
    return methodID;
}

bool JavaManager::callJavaFinish(JNIEnv *env) {
    if(g_dsManagerInstance == nullptr || g_dsFinish == nullptr){
        LOGE("DsManager instance or onFinsh is null\n");
        return false;
    }
//    5.2 调用 onFinsh 方法
    env->CallVoidMethod(g_dsManagerInstance, g_dsFinish);
    return true;
}


